Skip to content

Conversation

@thurstond
Copy link
Contributor

These functions will be useful in other intrinsic handlers.

thurstond added 2 commits July 7, 2025 22:17
…Zeros' into helper functions

These functions will be useful in other intrinsic handlers.
@llvmbot
Copy link
Member

llvmbot commented Jul 7, 2025

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-compiler-rt-sanitizer

Author: Thurston Dang (thurstond)

Changes

These functions will be useful in other intrinsic handlers.


Full diff: https://github.com/llvm/llvm-project/pull/147415.diff

1 Files Affected:

  • (modified) llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp (+54-20)
diff --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
index ec94dcaa2c051..566264f23c831 100644
--- a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
@@ -3389,6 +3389,58 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
     setOriginForNaryOp(I);
   }
 
+  /// Some instructions have additional zero-elements in the return type
+  /// e.g., <16 x i8> @llvm.x86.avx512.mask.pmov.qb.512(<8 x i64>, ...)
+  ///
+  /// This function will return a vector type with the same number of elements
+  /// as the input, but same width as the return value e.g., <8 x i8>.
+  FixedVectorType *shrinkVectorShadowType(Value *Src, IntrinsicInst &I) {
+    assert(isa<FixedVectorType>(getShadowTy(&I)));
+    // The return type might have more elements than the input.
+    // Temporarily shrink the return type's number of elements.
+    FixedVectorType *ShadowType = cast<FixedVectorType>(getShadowTy(&I));
+
+    // TODO: generalize beyond 2x?
+    if (ShadowType->getElementCount() ==
+        cast<VectorType>(Src->getType())->getElementCount() * 2)
+      ShadowType = FixedVectorType::getHalfElementsVectorType(ShadowType);
+
+    assert(ShadowType->getElementCount() ==
+           cast<VectorType>(Src->getType())->getElementCount());
+
+    return ShadowType;
+  }
+
+  /// Doubles the length of a vector shadow (filled with zeros) if necessary to
+  /// match the length of the shadow for the instruction.
+  /// This is more type-safe than CreateShadowCast().
+  Value *extendVectorShadowWithZeros(Value *Shadow, IntrinsicInst &I) {
+    IRBuilder<> IRB(&I);
+    assert(isa<FixedVectorType>(Shadow->getType()));
+    assert(isa<FixedVectorType>(I.getType()));
+
+    Value *FullShadow = getCleanShadow(&I);
+    assert(cast<FixedVectorType>(Shadow->getType())->getNumElements() <=
+           cast<FixedVectorType>(FullShadow->getType())->getNumElements());
+    assert(cast<FixedVectorType>(Shadow->getType())->getScalarType() ==
+           cast<FixedVectorType>(FullShadow->getType())->getScalarType());
+
+    if (Shadow->getType() == FullShadow->getType()) {
+      FullShadow = Shadow;
+    } else {
+      // TODO: generalize beyond 2x?
+      SmallVector<int, 32> ShadowMask(
+          cast<FixedVectorType>(FullShadow->getType())->getNumElements());
+      std::iota(ShadowMask.begin(), ShadowMask.end(), 0);
+
+      // Append zeros
+      FullShadow =
+          IRB.CreateShuffleVector(Shadow, getCleanShadow(Shadow), ShadowMask);
+    }
+
+    return FullShadow;
+  }
+
   /// Handle x86 SSE vector conversion.
   ///
   /// e.g., single-precision to half-precision conversion:
@@ -3419,13 +3471,7 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
 
     // The return type might have more elements than the input.
     // Temporarily shrink the return type's number of elements.
-    VectorType *ShadowType = cast<VectorType>(getShadowTy(&I));
-    if (ShadowType->getElementCount() ==
-        cast<VectorType>(Src->getType())->getElementCount() * 2)
-      ShadowType = VectorType::getHalfElementsVectorType(ShadowType);
-
-    assert(ShadowType->getElementCount() ==
-           cast<VectorType>(Src->getType())->getElementCount());
+    VectorType *ShadowType = shrinkVectorShadowType(Src, I);
 
     IRBuilder<> IRB(&I);
     Value *S0 = getShadow(&I, 0);
@@ -3440,19 +3486,7 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
 
     // The return type might have more elements than the input.
     // Extend the return type back to its original width if necessary.
-    Value *FullShadow = getCleanShadow(&I);
-
-    if (Shadow->getType() == FullShadow->getType()) {
-      FullShadow = Shadow;
-    } else {
-      SmallVector<int, 8> ShadowMask(
-          cast<FixedVectorType>(FullShadow->getType())->getNumElements());
-      std::iota(ShadowMask.begin(), ShadowMask.end(), 0);
-
-      // Append zeros
-      FullShadow =
-          IRB.CreateShuffleVector(Shadow, getCleanShadow(Shadow), ShadowMask);
-    }
+    Value *FullShadow = extendVectorShadowWithZeros(Shadow, I);
 
     setShadow(&I, FullShadow);
     setOriginForNaryOp(I);

assert(isa<FixedVectorType>(getShadowTy(&I)));
FixedVectorType *ShadowType = cast<FixedVectorType>(getShadowTy(&I));

// TODO: generalize beyond 2x?
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

shouldn't we fail the else case?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The assertion below will catch it in debug builds (though assuming sufficient test coverage, the else case or assertion will never happen).

Copy link
Contributor Author

@thurstond thurstond Jul 7, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

edit: no, we can't unconditionally fail the else case - sometimes it is already the correct number of elements.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybeShrinkVectorShadowType then?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done, maybe

Copy link
Contributor

@fmayer fmayer left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, probably

@thurstond thurstond changed the title [NFCI][msan] Extract 'shrinkVectorShadow' and 'extendVectorShadowWithZeros' into helper functions [NFCI][msan] Extract 'maybeShrinkVectorShadow' and 'maybeExtendVectorShadowWithZeros' into helper functions Jul 7, 2025
@thurstond thurstond merged commit 3528e16 into llvm:main Jul 8, 2025
9 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants